package com.terry.proxy.plugin.socks5;

import com.terry.proxy.util.ThreadPoolExecutorUtil;
import io.netty.bootstrap.Bootstrap;
import io.netty.bootstrap.ServerBootstrap;
import io.netty.channel.*;
import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.channel.socket.SocketChannel;
import io.netty.channel.socket.nio.NioServerSocketChannel;
import io.netty.channel.socket.nio.NioSocketChannel;
import io.netty.handler.codec.socksx.v5.*;
import io.netty.util.ReferenceCountUtil;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.boot.context.event.ApplicationReadyEvent;
import org.springframework.context.event.EventListener;
import org.springframework.stereotype.Component;

/**
 * socks5 插件
 *
 * @author terry
 * @version 1.0
 * @date 2023/2/3 11:16
 */
@Slf4j
@ConditionalOnProperty(name = "proxy.socks5.enable", havingValue = "true")
@Component
public class Socks5Plugin {

    @Value("${proxy.socks5.port}")
    private Integer port;

    @Value("${proxy.socks5.auth}")
    private boolean auth;

    @Value("${proxy.socks5.userName}")
    private String userName;

    @Value("${proxy.socks5.password}")
    private String password;

    @EventListener(ApplicationReadyEvent.class)
    public void start() {
        ThreadPoolExecutorUtil.execute(() -> Socks5Plugin.open(port, auth, userName, password));
    }

    /**
     * 开启 socks5
     * @param port socks5端口
     */
    public static void open(Integer port, boolean auth, String userName, String password){
        EventLoopGroup bossGroup = new NioEventLoopGroup();
        EventLoopGroup workerGroup = new NioEventLoopGroup();
        EventLoopGroup clientWorkGroup = new NioEventLoopGroup();
        try {
            ServerBootstrap b = new ServerBootstrap();
            b.group(bossGroup, workerGroup).channel(NioServerSocketChannel.class)
            .childHandler(new ChannelInitializer<SocketChannel>() {
                @Override
                public void initChannel(SocketChannel ch) {
                    ChannelPipeline pipeline = ch.pipeline();
                    // socks5响应最后一个encode
                    pipeline.addLast(Socks5ServerEncoder.DEFAULT);
                    // 处理socks5初始化请求
                    pipeline.addLast(new Socks5InitialRequestDecoder());
                    pipeline.addLast(new SimpleChannelInboundHandler<DefaultSocks5InitialRequest>() {
                        @Override
                        protected void channelRead0(ChannelHandlerContext ctx, DefaultSocks5InitialRequest msg) {
                            boolean failure = msg.decoderResult().isFailure();
                            if (failure) {
                                ReferenceCountUtil.retain(msg);
                                ctx.fireChannelRead(msg);
                                return;
                            }
                            if(auth){
                                Socks5InitialResponse socks5InitialResponse = new DefaultSocks5InitialResponse(Socks5AuthMethod.PASSWORD);
                                ctx.writeAndFlush(socks5InitialResponse);
                            }else{
                                Socks5InitialResponse socks5InitialResponse = new DefaultSocks5InitialResponse(Socks5AuthMethod.NO_AUTH);
                                ctx.writeAndFlush(socks5InitialResponse);
                            }
                            ctx.pipeline().remove(this);
                            ctx.pipeline().remove(Socks5InitialRequestDecoder.class);
                        }
                    });
                    // 处理认证请求
                    if (auth) {
                        pipeline.addLast(new Socks5PasswordAuthRequestDecoder());
                        pipeline.addLast(new Socks5AuthHandler(userName, password));
                    }
                    // 处理connection请求
                    pipeline.addLast(new Socks5CommandRequestDecoder());
                    pipeline.addLast(new Socks5Handler(clientWorkGroup));
                }
            });
            ChannelFuture channelFuture = b.bind(port).sync().channel().closeFuture().addListener((ChannelFutureListener) future -> {
                bossGroup.shutdownGracefully();
                workerGroup.shutdownGracefully();
            });
            log.info("Socks5插件已开启，占用端口234 {}", port);
            channelFuture.channel().closeFuture().sync();
        } catch (InterruptedException e) {
            e.printStackTrace();
        } finally {
            bossGroup.shutdownGracefully();
            workerGroup.shutdownGracefully();
        }
    }

    /**
     * socks5 处理类
     * @version 1.0
     * @author terry
     * @date 2022/12/29
     */
    @Slf4j
    @AllArgsConstructor
    public static class Socks5Handler extends SimpleChannelInboundHandler<DefaultSocks5CommandRequest> {

        private EventLoopGroup eventExecutors;

        @Override
        protected void channelRead0(ChannelHandlerContext ctx, DefaultSocks5CommandRequest msg) throws Exception {
            Socks5AddressType socks5AddressType = msg.dstAddrType();
            if (!msg.type().equals(Socks5CommandType.CONNECT)) {
                ReferenceCountUtil.retain(msg);
                ctx.fireChannelRead(msg);
                return;
            }
            log.info("连接目标服务器，ip：{},port：{}", msg.dstAddr(), msg.dstPort());
            Bootstrap bootstrap = new Bootstrap();
            bootstrap.group(eventExecutors)
            .channel(NioSocketChannel.class)
            .option(ChannelOption.TCP_NODELAY, true)
            .handler(new ChannelInitializer<SocketChannel>() {
                // 添加服务端写客户端
                @Override
                protected void initChannel(SocketChannel ch) throws Exception {
                    ch.pipeline().addLast(new ChannelInboundHandlerAdapter(){
                        @Override
                        public void channelRead(ChannelHandlerContext ctx2, Object msg) throws Exception {
                            ctx.writeAndFlush(msg);
                        }
                        @Override
                        public void channelInactive(ChannelHandlerContext ctx2) throws Exception {
                            ctx.channel().close();
                        }
                    });
                }
            });
            ChannelFuture future = bootstrap.connect(msg.dstAddr(), msg.dstPort());
            // 客户端转发请求到服务端
            future.addListener((ChannelFutureListener) future1 -> {
                if (future1.isSuccess()) {
                    ctx.pipeline().addLast(new ChannelInboundHandlerAdapter(){
                        @Override
                        public void channelRead(ChannelHandlerContext channelHandlerContext, Object msg1) throws Exception {
                            future1.channel().writeAndFlush(msg1);
                        }
                    });
                    DefaultSocks5CommandResponse commandResponse = new DefaultSocks5CommandResponse(Socks5CommandStatus.SUCCESS, socks5AddressType);
                    ctx.writeAndFlush(commandResponse);
                    ctx.pipeline().remove(Socks5Handler.class);
                    ctx.pipeline().remove(Socks5CommandRequestDecoder.class);
                } else {
                    log.error("连接目标服务器失败,ip：{},port：{}", msg.dstAddr(), msg.dstPort());
                    DefaultSocks5CommandResponse commandResponse = new DefaultSocks5CommandResponse(Socks5CommandStatus.FAILURE, socks5AddressType);
                    ctx.writeAndFlush(commandResponse);
                    future1.channel().close();
                }
            });
        }
    }
}
